#include <sys/types.h>
#include <sys/stat.h>
#include <sys/mman.h>
#include <unistd.h>
#include <stdlib.h>
#include <fcntl.h>

#include <um.h>
#include <pmm.h>
#include <mmu.h>
#include <memlayout.h>
#include <buddy_pmm.h>
#include <stdio.h>
#include <sync.h>
#include <types.h>
#include <stub.h>
#include <devices.h>

uintptr_t KERNBASE, KMEMSIZE = 0x08000000, VPT;

// virtual address of physicall page array
struct Page *pages;
// amount of physical memory (in pages)
size_t npage = 0;

// virtual address of boot-time page directory
pde_t *boot_pgdir = NULL;

// physical memory management
const struct pmm_manager *pmm_manager;


/**************************************************
 * Memory tests
 **************************************************/

/**
 * Check whether the virtual physical memory created works as designed.
 */
static void
check_vpm (void) 
{
	int* p;
	
	/* write at the beginning */
	p = (int*)KERNBASE;
	*p = 55;
	assert (*p == 55);

	/* write at the middle*/
	*p = 1111;
	assert (*p == 1111);

	/* write at the end */
	p = (int*)(UM_RESERVED - sizeof(int));
	*p = 101;
	assert (*p = 101);

	cprintf ("check_vpm() succeeded.\n");
	
	/* the code below will raise a segmentation fault */
	/* p = (int*)KERNTOP; */
	/* *p = 1000; */
}

/**************************************************
 * Memory manager wrappers.
 **************************************************/

/**
 * Initialize mm tragedy.
 */
static void
init_pmm_manager(void) {
	pmm_manager = &buddy_pmm_manager;
	cprintf("memory managment: %s\n", pmm_manager->name);
	pmm_manager->init();
}

/**
 * Initialize memory bitmap.
 */
static void
init_memmap(struct Page *base, size_t n) {
	pmm_manager->init_memmap(base, n);
}

/**
 * Check the correctness of the system.
 */
static void
check_alloc_page(void) {
	pmm_manager->check();
	cprintf("check_alloc_page() succeeded.\n");
}

/**************************************************
 * Page operation wrappers.
 **************************************************/
struct Page *
alloc_pages(size_t n) {
	struct Page *page;
	bool intr_flag;
	local_intr_save(intr_flag);
	{
		page = pmm_manager->alloc_pages(n);
	}
	local_intr_restore(intr_flag);
	return page;
}


void
free_pages(struct Page *base, size_t n) {
	bool intr_flag;
	local_intr_save(intr_flag);
	{
		pmm_manager->free_pages(base, n);
	}
	local_intr_restore(intr_flag);
}


size_t
nr_free_pages(void) {
	size_t ret;
	bool intr_flag;
	local_intr_save(intr_flag);
	{
		ret = pmm_manager->nr_free_pages();
	}
	local_intr_restore(intr_flag);
	return ret;
}

/**************************************************
 * Page operation wrappers end.
 **************************************************/

/**
 * Create virtual physical memory in /tmp
 * @param size the size of memory
 */
static int
create_mem_file (uint32_t size)
{
	int fd;
	char tempname[] = "/tmp/vpmXXXXXX";
	
	fd = mkstemp (tempname);
	if (fd < 0) {
		cperror ("create temp file failed.\n");
		_exit (1);
	}
	if (unlink (tempname) < 0) {
		cperror ("unlink failed\n");
		_exit (1);
	}
	
	/* Expand the file to the desired length */
	cprintf ("memory size: 0x%08x\n", size);
	if (lseek (fd, size - 1, SEEK_SET) < 0) {
		cperror ("lseek failed\n");
		_exit (1);
	}
	char zero = 0;
	if (write (fd, &zero, 1) < 0) {
		cperror ("write failed\n");
		_exit (1);
	}

	/* Change permission of the file */
	if (fchmod (fd, 0777) < 0) {
		cperror ("fchmod failed\n");
		_exit (1);
	}
	fcntl (fd, F_SETFD, FD_CLOEXEC);

	return fd;
}

static struct mmap_arg_struct mmap_args;

/**
 * Examine address space, create virtual physical memory and map it.
 */
static void
page_init (void)
{
	int i;
	
	uintptr_t brk_start = (uintptr_t) sbrk (0);
	KERNBASE = ROUNDUP (brk_start, PGSIZE);

	device_common->vpm_fd = create_mem_file (KMEMSIZE);
	mmap_args.addr = KERNBASE;
	mmap_args.len = KMEMSIZE - UM_RESERVED_SIZE;
	mmap_args.prot = PROT_EXEC|PROT_READ|PROT_WRITE;
	mmap_args.flags = MAP_SHARED|MAP_FIXED;
	mmap_args.fd = device_common->vpm_fd;
	mmap_args.offset = 0;
	void* loc = (void*)stub_syscall1 (__NR_mmap, (int)&mmap_args);
	if ((int)loc != KERNBASE) {
		cperror ("mmap failed\n");
		_exit (1);
	}
	cprintf ("vpm mapped at [0x%08x, 0x%08x]\n", KERNBASE, KERNTOP-1);
	check_vpm ();
	
	/* Construct page descriptor table. */
	pages = (struct Page *)(KERNBASE);
	npage = (UM_RESERVED - KERNBASE) / PGSIZE;
	for (i = 0; i < npage; i++) {
		SetPageReserved (pages + i);
	}
	uintptr_t freemem = PADDR(ROUNDUP((uintptr_t)pages + sizeof(struct Page) * npage, PGSIZE));
	uint32_t freemem_npage = npage - ROUNDUP(sizeof(struct Page) * npage, PGSIZE) / PGSIZE;
	init_memmap(pa2page(freemem), freemem_npage);
}


/**
 * map necessary pages to the address space
 */
static void
map_memory (void)
{
	cprintf ("syscall_stub_start = 0x%x\tstart = 0x%x\n", (int)&__syscall_stub_start, (int)stub_start);
	/* Write the stubs to the lowest in the reserved area. */
	lseek (device_common->vpm_fd, PADDR(UM_RESERVED), SEEK_SET);
	write (device_common->vpm_fd, &__syscall_stub_start, PGSIZE);
	
	/* Map the stub area to UM_STUB_CODE */
	mmap_args.addr = UM_STUB_CODE;
	mmap_args.len = PGSIZE;
	mmap_args.flags = MAP_SHARED|MAP_FIXED;
	mmap_args.prot = PROT_EXEC|PROT_READ;
	mmap_args.offset = PADDR(UM_RESERVED);
	stub_syscall1 (__NR_mmap, (int)&mmap_args);
	
	/* Map UM_STUB_DATA */
	mmap_args.addr = UM_STUB_DATA;
	mmap_args.len = PGSIZE;
	mmap_args.flags = MAP_SHARED|MAP_FIXED;
	mmap_args.prot = PROT_WRITE|PROT_READ;
	mmap_args.offset = PADDR(UM_RESERVED) + PGSIZE;
	stub_syscall1 (__NR_mmap, (int)&mmap_args);
	/* Allocate a page for the first child process as its stack */
	struct Page* stack_page = alloc_page ();
	mmap_args.addr = UM_SPACE_TOP - PGSIZE;
	mmap_args.len = PGSIZE;
	mmap_args.flags = MAP_SHARED|MAP_FIXED;
	mmap_args.prot = PROT_WRITE|PROT_READ;
	mmap_args.offset = page2pa (stack_page);
	lseek (device_common->vpm_fd, PADDR(UM_RESERVED) + PGSIZE, SEEK_SET);
	write (device_common->vpm_fd, &mmap_args, sizeof (struct mmap_arg_struct));
}


/**
 * Create tmp-file-based virtual physical memory,
 *    map it to the address space of the current process (tracing thread)
 *    and build buddy system on it.
 */
void
pmm_init (void) 
{
	init_pmm_manager ();

	/* Create memory and its free page list. */
	page_init ();

	check_alloc_page ();

	map_memory ();
}
